import torch
from loader import load
from torch.optim import AdamW
from tqdm import tqdm
from prim_net import prim_net as prim
from moe_net import moe_net as moe
from moe_net import head_counter


if __name__=='__main__':
    net=prim()
    NUM_EPOCH=2
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    optimizer=AdamW(params=net.parameters())
    train_loader, test_loader, len_train, len_test=load()
    net.to(device)
    loss_fn=torch.nn.CrossEntropyLoss()
    
    # train
    for epoch in range(NUM_EPOCH):
        total_train_loss=0
        total_test_loss=0
        train_correct=0
        test_correct=0
        
        net.train()
        for i,(inputs,labels) in enumerate(tqdm(train_loader)):
            inputs.to(device)
            labels.to(device)
            
            
            result=net(inputs)
            loss=loss_fn(result,labels)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            
            total_train_loss+=loss.item()

            _, predicted = torch.max(result.data, 1)  # 获取预测结果
            train_correct += (predicted == labels).sum().item()  # 统计预测正确

        net.eval()
        for i,(inputs,labels) in enumerate(tqdm(test_loader)):
            inputs.to(device)
            labels.to(device)
            
            
            #result=net(inputs,count=True)   #专家头计数
            result=net(inputs)
            loss=loss_fn(result,labels)
            
            total_test_loss+=loss.item()
            
            _, predicted = torch.max(result.data, 1)  # 获取预测结果
            test_correct += (predicted == labels).sum().item()  # 统计预测正确
            
        
        print(f"train_loss:{total_train_loss/len_train} train_correct:{train_correct/len_train}")
        print(f"test_loss:{total_test_loss/len_test} test_correct:{test_correct/len_test}")
        #print(head_counter)        #专家头的利用率